# -*- coding: utf-8 -*-  
'''
零件
    CRFLayer
        1 转移矩阵定义
        2 fc层定义
        3 返回fc层结果
            - crf返回结果与tplinker无法保存进1个tensor中
    
    
@author: luoyi
Created on 2021年7月4日
'''
import tensorflow as tf

import utils.conf as conf
import utils.poses as poses
import utils.relationships as rel


#    crf层
class CRFLayer(tf.keras.layers.Layer):
    def __init__(self,
                 name='crf_layer',
                 pos_size=poses.bio_tag_size(),
                 max_sen_len=conf.CRFLINKER.get_max_sentence_len(),
                 **kwargs):
        super(CRFLayer, self).__init__(name=name, **kwargs)
        
        self._pos_size = pos_size
        self._max_sen_len = max_sen_len
        
        pass
    
    def build(self, input_shape):
        #    状态转移矩阵
        self._transfer = self.add_weight(name='crf_transfer', 
                                         shape=(self._pos_size, self._pos_size), 
                                         dtype=tf.float32, 
                                         initializer=tf.initializers.he_normal(), 
                                         trainable=True)
        
        #    fc层
        self._fc = tf.keras.layers.Dense(name='crf_fc',
                                         units=self._pos_size, 
                                         kernel_initializer=tf.initializers.he_normal(), 
                                         bias_initializer=tf.initializers.he_normal())
        #    fc层输出
        self._out = None
        
        super(CRFLayer, self).build(input_shape)
        pass
    
    def call(self, inputs, **kwargs):
        #    inputs为bert输出 Tensor(batch_size, max_sen_len, d_model)
        #    过fc层，拿到crf输出
        out = self._fc(inputs)
        self._out = out
        
        return out
    
    #    拿fc层输出
    def get_out(self):
        return self._out
    def get_transfer(self):
        return self._transfer
    
    pass


#    tplinker层
class TPLinkerLayer(tf.keras.layers.Layer):
    def __init__(self,
                 name='tplinker_layer',
                 rel_size=len(rel.id_rel),
                 max_sen_len=conf.CRFLINKER.get_max_sentence_len(),
                 d_model=conf.BERT.get_d_model(),
                 **kwargs):
        super(TPLinkerLayer, self).__init__(name=name, **kwargs)
        
        self._rel_size = rel_size
        self._max_sen_len = max_sen_len
        self._d_model = d_model
        
        #    向量长度
        self._max_vec_len = ((1 + max_sen_len-2) * (max_sen_len-2)) // 2
        pass
    
    def build(self, input_shape):
        #    握手层
        self._handshaking_layer = HandshakingLayer(rel_size=self._rel_size, 
                                                   max_sen_len=self._max_sen_len,
                                                   d_model=self._d_model)
        
        #    fc层
        self._fc = tf.keras.layers.Dense(name='tplinker_fc', 
                                         units=self._rel_size * 2 * 3, 
                                         kernel_initializer=tf.initializers.he_normal(), 
                                         bias_initializer=tf.initializers.he_normal(), 
                                         )
        
        #    fc层结果
        self._out = None
        
        super(TPLinkerLayer, self).build(input_shape)
        pass
    
    def call(self, inputs, **kwargs):
        #    inputs为bert输出 Tensor(batch_size, max_sen_len, d_model)
        vec = self._handshaking_layer(inputs)
        
        #    过fc层
        vec = self._fc(vec)
        vec = tf.reshape(vec, shape=(-1, self._max_vec_len, self._rel_size, 2, 3))
        vec = tf.transpose(vec, perm=[0, 2, 1, 3, 4])                   #    Tensor(batch_size, rel_size, max_vec_len, 2, 3)
        #    3个预测值过softmax
        vec = tf.nn.softmax(vec, axis=-1)
        
        #    记录本层运算结果
        self._out = vec
        
        return vec
    
    def get_out(self):
        return self._out
        
    pass


#    握手层
class HandshakingLayer(tf.keras.layers.Layer):
    def __init__(self,
                 name='handshaking_layer',
                 rel_size=len(rel.id_rel),
                 max_sen_len=conf.CRFLINKER.get_max_sentence_len(),
                 d_model=conf.BERT.get_d_model(),
                 **kwargs):
        super(HandshakingLayer, self).__init__(name=name, **kwargs)
        
        self._rel_size = rel_size
        self._max_sen_len = max_sen_len
        self._d_model = d_model
        
        #    向量长度
        self._max_vec_len = ((1 + max_sen_len-2) * (max_sen_len-2)) // 2
        pass
    
    def build(self, input_shape):
        self._fc = tf.keras.layers.Dense(name=self.name + '_fc',
                                         units=self._d_model, 
                                         kernel_initializer=tf.initializers.he_normal(), 
                                         bias_initializer=tf.initializers.he_normal())
        
        super(HandshakingLayer, self).build(input_shape)
        pass
    
    def call(self, inputs, **kwargs):
        #    inputs为bert输出 Tensor(batch_size, max_sen_len, d_model)
        x = inputs[:, 1:-1, :]                      #    去掉每行的[CLS][EOS]
                                                    #    [CLS]位置是对的，[EOS]可能不对。但无所谓，下面是通过索引取的
        (_, sen_len, d_model) = x.shape    #    sen_len = max_sen_len - 2
                                                    
        x_i = tf.tile(x[:, :, tf.newaxis, :], multiples=[1, 1, sen_len, 1])        #    x_i[i,j] = x[i]
        x_j = tf.tile(x[:, tf.newaxis, :, :], multiples=[1, sen_len, 1, 1])        #    x_j[i,j] = x[j]
        x = tf.concat([x_i, x_j], axis=-1)                                          #    x[i,j] = [x[i] ; x[j]]
        
        #    过fc层，tanh激活
        x = self._fc(x)
        x = tf.nn.tanh(x)
        
        #    取上三角部分（包括对角线）
        mat = tf.ones_like(x)[:, :, :, 0]                                           #    Tensor(batch_size, sen_len, sen_len)
        mat = tf.linalg.band_part(mat, num_lower=0, num_upper=sen_len)              #    Tensor(batch_size, sen_len, sen_len)，下三角部分为0
        mat = tf.cast(mat, dtype=tf.int64)
        x_vec = tf.gather_nd(x, indices=tf.where(mat == 1))
        x_vec = tf.reshape(x_vec, shape=(-1, self._max_vec_len, d_model))
        
        return x_vec
        
    pass
